package web

import (
	"gitee.com/crack007/goose/gsutil"
	"github.com/gin-gonic/gin"
)

var JsonMiddleware = jsonMiddleware{}

type jsonMiddleware struct {
}

// 自定义的 ResponseWriter，用于捕获响应数据
type responseWriter struct {
	gin.ResponseWriter
	body []byte
}

// 重写 Write 方法，捕获响应数据。
func (w *responseWriter) Write(b []byte) (int, error) {
	w.body = append(w.body, b...) // 捕获响应数据
	return w.ResponseWriter.Write(b)
}

func (x *jsonMiddleware) Apply(iRoutes gin.IRoutes) {
	iRoutes.Use(x.Handle)
}

func (x *jsonMiddleware) Handle(ctx *gin.Context) {
	if ctx.Keys == nil {
		ctx.Keys = map[string]any{}
	}

	ctx.Next()
	dto := GetResponse(ctx)
	var resp string
	if v, ok := dto.(*ApiException); ok {
		resp = gsutil.JsonUtil.Marshal(gin.H{
			"msg":  v.Msg,
			"code": v.Code,
			"data": dto,
		})
	} else {
		resp = gsutil.JsonUtil.Marshal(gin.H{
			"msg":  "",
			"code": 200,
			"data": dto,
		})
	}

	writer := &responseWriter{ResponseWriter: ctx.Writer, body: []byte{}}
	ctx.Writer = writer

	ctx.Header("Content-Type", "application/json")
	ctx.Writer.Write([]byte(resp))
}
